module TrainManagement where 

import DataTypes
import DataStorage
import GenericFunctions
import StationManagement
import Data.List
import Data.Maybe
import Data.Time.Format
import Data.Time.Calendar
import Data.Time.LocalTime
import Control.Monad
import System.IO
import System.Locale


trainsFile = "trains"

-- *******************************
-- * Trains management interface *
-- *******************************

-- Function used to print out all trains
printTrains  = getTrains >>= printTrains'
                 
printTrains' trlist =   do  str <- mapM (\t -> trainStr t) trlist
                            showMenu "Available trains" $ (zip str [printTrainStations t >> pressEnter| t <- trlist]) ++ [("CANCEL", return ())]

-- Function used to read information about a train
promptTrain = getTrains >>= promptTrain'

promptTrain' trlist =   do  str <- mapM (\t -> trainStr t) trlist
                            showMenu "Train choice" $ (zip str [return (Just t) | t <- trlist]) ++ [("[CANCEL]", return Nothing)]

-- Function used to print all stations of a specified train                              
printTrainStations train =  do  (r1, r2) <- getTrainRoutes train
                                showItems ("Train \"" ++ name train ++ "\" stations [direction 1]") [stationStr s | s <- r1]
                                showItems' "schedule [direction 1]" (fst $ trainScheduleItems train) True
                                showItems ("Train \"" ++ name train ++ "\" stations [direction 2]") [stationStr s | s <- r2]
                                showItems' "schedule [direction 2]" (snd $ trainScheduleItems train) True
                                
trainScheduleItems train = ((concat [" station " ++ (show n) ++ "  | " | n <- [1..(length $ fst (route train))]]) : [schedStr $ fst sch | sch <- sched train],
                            (concat [" station " ++ (show n) ++ "  | " | n <- [1..(length $ snd (route train))]]) : [schedStr $ snd sch | sch <- sched train])
                                where
                                    schedStr s = concat [(timeStr $ arrival di) ++ "/" ++ (timeStr $ departure di) ++ " | "| di <- s]

-- Function used to get all train depatures from a given station
departuresFromStation train sid =   (deps train sid fst, deps train sid snd)
                                    where   deps t s f = [di | di <- concat [f d | d <- sched t], stationId di == s]
                
-- Function used to add a new train
addTrain =  do  tname <- promptLine "New train name (or old name to add a new course to an existing train)"
                train <- getTrainByName tname
                if (isJust train) then do                    
                    answer <- promptLine "The train with given name already exists. Would you like to add a new course? (y/n)"
                    if (answer == "y") then do
                        let (Train name since till r s) = fromJust train
                        (r1, r2) <- getTrainRoutes (fromJust train)
                        let n1 = stationName (head r1)
                        let n2 = stationName (last r1)
                        let n3 = stationName (last r2)
                        putStrLn ("Loading timetable of relation \"" ++ n1 ++ "\" -> \"" ++ n2 ++ "\"")
                        sched1 <- getDepatures r1 Forward
                        putStrLn ("Loading timetable of relation \"" ++ n2 ++ "\" -> \"" ++ n3 ++ "\"")
                        sched2 <- getDepatures r2 Backward                       
                        getTrains >>= deleteTrain (fromJust train)
                        saveTrain (Train tname since till r (s ++ [(sched1, sched2)]))
                        return ()
                    else
                        return ()
                else do
                    putStrLn ("Adding a new train " ++ tname)
                    sinceDate <- promptDate "Runs since"
                    tillDate <- promptDate "Runs till"
                    if (tillDate > sinceDate) then do
                        stations <- StationManagement.getStations
                        route1 <- loadRoute stations [] "Selected route \"there\""
                        let r2start = last route1
                        let stations2 = filter (\x -> x /= r2start) stations
                        route2 <- loadRoute stations2 [r2start] "Selected route \"back\""
                        let (Station _ n1 _) = head route1
                        let (Station _ n2 _) = last route1
                        let (Station _ n3 _) = last route2
                        putStrLn ("Loading timetable of relation \"" ++ n1 ++ "\" -> \"" ++ n2 ++ "\"")
                        sched1 <- getDepatures route1 Forward
                        putStrLn ("Loading timetable of relation \"" ++ n2 ++ "\" -> \"" ++ n3 ++ "\"")
                        sched2 <- getDepatures route2 Backward
                        let route1s = serializeRoute route1
                        let route2s = serializeRoute route2
                        saveTrain (Train tname sinceDate tillDate (route1s, route2s) [(sched1, sched2)])
                        return ()
                    else
                        putStrLn "error: till before since"

-- Function used to load a train route
loadRoute [] a _  = return a
loadRoute slist clist title =   do  showItems title $ map (\(Station _ name _) -> ("\"" ++ name ++ "\" ")) clist                            
                                    mstation <- StationManagement.promptStation' slist
                                    if (isJust mstation) then do
                                        let station = fromJust mstation
                                        let slist2 = filter (\(Station id _ _) -> id /= DataTypes.id station) slist                                
                                        x <- loadRoute slist2 (clist ++ [station]) title
                                        return x
                                    else
                                        return clist

-- Function used to edit an existing train
editTrain = do  trains <- getTrains
                mtrain <- promptTrain' trains
                if (isJust mtrain) then do
                    let train = fromJust mtrain        
                    let rest = filter (\t -> name t /= name train) trains
                    DataStorage.saveItems rest trainsFile
                    changeTrain train
                else return ()
                    
-- Function used to remove a train
removeTrain =   do  trains <- getTrains                    
                    mtrain <- promptTrain' trains
                    when (isJust mtrain) (deleteTrain (fromJust mtrain) trains)
                    
-- Function used to save a train
saveTrain train = DataStorage.addItem train trainsFile   
                            
-- Function used to remove a train
deleteTrain t trains  = DataStorage.saveItems (filter (\tx -> name t /= name tx) trains) trainsFile
                    
-- Function used to get all trains
getTrains = DataStorage.loadItems trainsFile
                
-- Function used to get a train with a given name
getTrainByName tname =  do  trains <- getTrains
                            return (find (\(Train name _ _ _ _) -> name == tname) trains)
                            
-- Function used to print a readable train description                            
trainStr (Train name since till (r1, r2) [(s1, s2)])  = 
    do
        st1 <- StationManagement.getStationById (head r1)
        st2 <- StationManagement.getStationById (last r1)
        st3 <- StationManagement.getStationById (last r2)
        let name1 = stationName $ fromJust st1
        let name2 = stationName $ fromJust st2
        let name3 = stationName $ fromJust st3
        return $ "Train \"" ++ name ++ "\", \"" ++ name1 ++ "\" -> \"" ++ name2 ++ "\" -> \"" ++ name3 ++ "\", [" ++ (dateStr since) ++ ", " ++ (dateStr till) ++ "]"

-- Function used to print a train
printTrain train = trainStr train >>= putStrLn
                             
-- Function used to get all depatures from a given station
getDepatures route dir = mapM (getDepartureItem dir) route
                            
-- Function used to get depature item from the given station                            
getDepartureItem dir (Station id name city) =   do  arrivalTime <- promptTime $ "Arrival at the station \"" ++ name ++ "\""
                                                    departsTime <- promptTime $ "Depature from the station \"" ++ name ++ "\""
                                                    return (DeparturesItem arrivalTime departsTime id dir)
         
-- get routes of a given train
getTrainRoutes train =  do  r1 <- deserializeRoute $ fst (route train)
                            r2 <- deserializeRoute $ snd (route train)
                            return (r1, r2)
                            
-- Serialize route of a train
serializeRoute route = map (\(Station id _ _) -> id) route

-- Deserialize route of a train
deserializeRoute route =    do  rm <- mapM (\rs -> StationManagement.getStationById rs) route
                                return (map (\r -> fromJust r) rm)

-- Change the train name
changeTrainName name = do answer <- promptLine "Do you want to change train name? (Y/N)"
                          if (answer == "Y") then do
                             tname <- promptLine "New train name"
                             train <- getTrainByName tname
                             if (tname /= name && (isJust train)) then do
                                error "Such train already exists!"
                             else do return tname
                          else do return name

-- Change the train route
changeTrainRoute route = do answer <- promptLine "Would you like to drop previous routes and add a new one? (Y/N)"
                            if (answer == "Y") then do
                               stations <- StationManagement.getStations
                               route1 <- loadRoute stations [] "Selected route \"there\""
                               let r2start = last route1
                               let stations2 = filter (\x -> x /= r2start) stations
                               route2 <- loadRoute stations2 [r2start] "Selected route \"back\""
                               let route1s = serializeRoute route1
                               let route2s = serializeRoute route2
                               return (route1s, route2s)
                            else do return route

-- Change the train date
changeTrainDate:: Day -> String -> IO Day
changeTrainDate date msg = do answer <- promptLine msg
                              if (answer == "Y") then do
                                  promptDate "New date"
                              else do return date

-- Change the train timetable
changeTrainTimetable table (route1s, route2s) = 
                    do route1 <- deserializeRoute route1s
                       route2 <- deserializeRoute route2s
                       let (Station _ n1 _) = head route1
                       let (Station _ n2 _) = last route1
                       let (Station _ n3 _) = last route2
                       putStrLn ("Loading timetable of relation \"" ++ n1 ++ "\" -> \"" ++ n2 ++ "\"")
                       sched1 <- getDepatures route1 Forward
                       putStrLn ("Loading timetable of relation \"" ++ n2 ++ "\" -> \"" ++ n3 ++ "\"")
                       sched2 <- getDepatures route2 Backward
                       return [(sched1, sched2)]

-- Change the train
changeTrain (Train name since till route sched) = 
              do tname <- changeTrainName name
                 troute <- changeTrainRoute route
                 tsince <- changeTrainDate since "Would you like to change since date (Y/N)?"
                 ttill <- changeTrainDate till "Would you like to change till date (Y/N)?"
                 if (troute /= route) then do
                     tsched <- changeTrainTimetable sched troute
                     saveTrain (Train tname tsince ttill troute tsched)
                 else do 
                    answer <- promptLine "Would you like to change timetable? (Y/N)"
                    if (answer == "Y") then do
                        tsched <- changeTrainTimetable sched troute
                        saveTrain (Train tname tsince ttill troute tsched)
                    else do saveTrain (Train tname tsince ttill troute sched)

